import os
import asyncio
import uuid
from datetime import datetime
from flask import Flask, request, jsonify, send_from_directory, render_template
from TTS.api import TTS
import edge_tts
from pydub import AudioSegment

app = Flask(__name__, static_folder='static', template_folder='templates')
app.config['UPLOAD_FOLDER'] = 'TTS/audio'
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# 优化模型加载
class TTSAudio:
    def __init__(self, language='en'):
        self.tts = {
            "zh": "tts_models/zh-CN/baker/tacotron2-DDC-GST",
            "en": "tts_models/en/ljspeech/tacotron2-DDC",
            "multi": "tts_models/multilingual/multi-dataset/xtts_v2"
        }
        self.language = language
        self.model = None

    def get_model(self, model_name):
        if self.model is None:  # 仅加载一次模型
            self.model = TTS(model_name=self.tts[model_name], progress_bar=False)
        return self.model


@app.route('/')
def index():
    return render_template('tts_web.html')  # 渲染前端页面


@app.route('/tts', methods=['POST'])
def text_to_speech():
    """文本转语音接口"""
    data = request.get_json()
    if not data or 'text' not in data:
        return jsonify({'error': '缺少 text 参数'}), 400

    text = data['text']
    if not text.endswith('。'):
        text += '。'

    language = data.get('language', 'en')
    model = data.get('model', 'en')
    speed = float(data.get('speed', 1.0))

    # 生成唯一音频文件名
    filename = f"tts_{datetime.now().strftime('%Y%m%d%H%M%S')}_{uuid.uuid4().hex[:6]}.wav"
    output_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)

    try:
        if model == 'multi':
            # XTTS v2 多语言模型
            tts_engine = TTSAudio(language)
            tts = tts_engine.get_model(model)
            tts.tts_to_file(
                text=text,
                language=language,
                speaker_wav="config/v1fbphs4.wav",
                speed=speed,
                file_path=output_path
            )
        elif model == 'yue':
            # 直接调用 Edge TTS（修正）
            asyncio.run(generate_edge_tts(text, output_path))
        else:
            # 普通 TTS 处理
            tts_engine = TTSAudio(language)
            tts = tts_engine.get_model(model)
            tts.tts_to_file(
                text=text,
                file_path=output_path
            )

        # 返回音频文件 URL
        return jsonify({
            'audio_url': f"/download/{filename}",
            'text': text,
            'params': {
                'language': language,
                'speed': speed
            }
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/download/<filename>')
def download_file(filename):
    """提供音频文件下载"""
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)


async def generate_edge_tts(text, output_path):
    """Edge TTS 生成音频，并转为 WAV"""
    voice = "zh-HK-HiuMaanNeural"
    temp_mp3 = output_path.replace(".wav", ".mp3")
    tts = edge_tts.Communicate(text, voice)
    await tts.save(temp_mp3)
    audio = AudioSegment.from_mp3(temp_mp3)
    audio.export(output_path, format="wav")


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)